# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Conformation embedding module for molecular representation
"""

from typing import Tuple, Optional, Union, Any

import torch
from torch import nn, Tensor

from ..utils import GLOBAL_DEVICE
from ..units import Length, GLOBAL_UNITS

class ConformationEmbedding(nn.Module):
    r"""Embedding module for molecular conformation information.

    Args:
        dim_feature (int): Dimension of the output feature vectors.
        emb_bond (bool, optional): Whether to include bond information in embedding.
            Default: False
        cutoff (float, optional): Cutoff distance for interactions.
            Default: Length(1, 'nm')
        cutoff_fn (Any, optional): Cutoff function to use.
            Default: None
        rbf_fn (Any, optional): Radial basis function to use.
            Default: None
        num_basis (int, optional): Number of basis functions.
            Default: None
        atom_filter (Any, optional): Filter for atom representations.
            Default: None
        dis_filter (str, optional): Filter type for distance features.
            Default: 'residual'
        bond_filter (str, optional): Filter type for bond features.
            Default: 'residual'
        interaction (Any, optional): Interaction module to use.
            Default: None
        dis_self (Any, optional): Self-distance calculation method.
            Default: Length(0.05, 'nm')
        num_atom_types (int, optional): Number of atom types.
            Default: 64
        num_bond_types (int, optional): Number of bond types.
            Default: 16
        initializer (Any, optional): Weight initializer.
            Default: None
        activation (str, optional): Activation function to use.
            Default: 'swish'
        length_unit (Any, optional): Unit for length measurements.
            Default: None
    """
    def __init__(self,
                 dim_feature: int,
                 emb_bond: bool = False,
                 cutoff: Length = Length(1, 'nm'),
                 cutoff_fn: Optional[Any] = None,
                 rbf_fn: Optional[Any] = None,
                 num_basis: Optional[int] = None,
                 atom_filter: Optional[Any] = None,
                 dis_filter: str = 'residual',
                 bond_filter: str = 'residual',
                 interaction: Optional[Any] = None,
                 dis_self: Length = Length(0.05, 'nm'),
                 num_atom_types: int = 64,
                 num_bond_types: int = 16,
                 initializer: Optional[Any] = None,
                 activation: str = 'swish',
                 length_unit: str = GLOBAL_UNITS.length_unit,
                 ):
        super().__init__()
        self.dim_feature = dim_feature
        self.emb_bond = emb_bond
        self.cutoff = cutoff
        self.cutoff_fn = cutoff_fn
        self.rbf_fn = rbf_fn
        self.num_basis = num_basis
        self.atom_filter = atom_filter
        self.dis_filter = dis_filter
        self.bond_filter = bond_filter
        self.interaction = interaction
        self.dis_self = dis_self
        self.num_atom_types = num_atom_types
        self.num_bond_types = num_bond_types
        self.initializer = initializer
        self.activation = activation
        self.length_unit = length_unit
        self.device = GLOBAL_DEVICE()

    def forward(self,
                atom_type: Tensor,
                atom_mask: Tensor,
                distance: Tensor,
                dis_mask: Tensor,
                bond: Tensor,
                bond_mask: Tensor,
                **kwargs,
                ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        return NotImplementedError
